import paddle
from paddle import nn
import paddle.nn.functional as F
from paddle import ParamAttr



def conv_bn_relu(in_channels, out_channels, kernel_size, stride):
    pad = kernel_size // 2
    return nn.Sequential(
        nn.Conv2D(in_channels, out_channels, kernel_size, stride, pad, bias_attr=False),
        nn.BatchNorm2D(out_channels),
        nn.ReLU()
    )


class DBUNet(nn.Layer):
    def __init__(self, in_channels, out_channels, **kwargs):
        super().__init__()
        self.out_channels = out_channels
        self.ups = nn.LayerList()
        self.convs = nn.LayerList()

        outplanes = out_channels
        num = len(in_channels)
        for i in range(num - 1):
            self.ups.append(self._make_upsample(in_channels[i+1], in_channels[i]))
            self.convs.append(conv_bn_relu(in_channels[i]*2, outplanes, 3, 1))
            outplanes = in_channels[i+1]


    def _make_upsample(self, in_channels, out_channels, scale=2.0):
        return nn.Sequential(
            nn.UpsamplingBilinear2D(scale_factor=scale),
            nn.Conv2D(in_channels, out_channels, 3, 1, 1, bias_attr=False),
            nn.BatchNorm2D(out_channels),
            nn.ReLU()
        )
    
    def forward(self, feat):
        num = len(feat)
        x = feat[-1]
        for i in range(num-1, 0, -1):
            x = self.ups[i-1](x)
            x = paddle.concat([x, feat[i-1]], 1)
            x = self.convs[i-1](x)

        return x
